// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <algorithm>
#include <limits>
#include <optional>

#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/sparksql/registration/register.h>
#include <pollux/functions/sparksql/tests/ArraySortTestData.h>
#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>
#include <pollux/vector/complex_vector.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions::sparksql::test {
namespace {

class SortArrayTest : public SparkFunctionBaseTest {
 protected:
  void testSortArray(
      const VectorPtr& input,
      const VectorPtr& expectedAsc,
      const VectorPtr& expectedDesc) {
    // Verify that by default array is sorted in ascending order.
    testSortArray("sort_array(c0)", input, expectedAsc);

    // Verify sort order with asc flag set to true.
    testSortArray("sort_array(c0, true)", input, expectedAsc);

    // Verify sort order with asc flag set to false.
    testSortArray("sort_array(c0, false)", input, expectedDesc);
  }

  void testSortArray(
      const std::string& expr,
      const VectorPtr& input,
      const VectorPtr& expected) {
    SCOPED_TRACE(expr);
    auto result = evaluate(expr, make_row_vector({input}));
    assertEqualVectors(expected, result);
  }

  template <typename T>
  void testInt() {
    auto input = make_nullable_array_vector(intInput<T>());
    auto expected = intAscNullSmallest<T>();
    testSortArray(
        input,
        make_nullable_array_vector(expected),
        make_nullable_array_vector(reverseNested(expected)));
  }

  template <typename T>
  void testFloatingPoint() {
    auto input = make_nullable_array_vector(floatingPointInput<T>());
    auto expected = floatingPointAscNullSmallest<T>();
    testSortArray(
        input,
        make_nullable_array_vector(expected),
        make_nullable_array_vector(reverseNested(expected)));
  }
};

TEST_F(SortArrayTest, invalidInput) {
  auto arg0 = make_nullable_array_vector<int>({{0, 1}});
  std::vector<bool> v = {false};
  auto arg1 = make_flat_vector<bool>(v);
  ASSERT_THROW(
      evaluate<ArrayVector>("sort_array(c0, c1)", make_row_vector({arg0, arg1})),
      PolluxException);
}

TEST_F(SortArrayTest, int8) {
  testInt<int8_t>();
}

TEST_F(SortArrayTest, int16) {
  testInt<int16_t>();
}

TEST_F(SortArrayTest, int32) {
  testInt<int32_t>();
}

TEST_F(SortArrayTest, int64) {
  testInt<int64_t>();
}

TEST_F(SortArrayTest, float) {
  testFloatingPoint<float>();
}

TEST_F(SortArrayTest, double) {
  testFloatingPoint<double>();
}

TEST_F(SortArrayTest, string) {
  auto input = make_nullable_array_vector(stringInput());
  auto expected = stringAscNullSmallest();
  testSortArray(
      input,
      make_nullable_array_vector(expected),
      make_nullable_array_vector(reverseNested(expected)));
}

TEST_F(SortArrayTest, timestamp) {
  auto input = make_nullable_array_vector(timestampInput());
  auto expected = timestampAscNullSmallest();
  testSortArray(
      input,
      make_nullable_array_vector(expected),
      make_nullable_array_vector(reverseNested(expected)));
}

TEST_F(SortArrayTest, date) {
  auto input = make_nullable_array_vector(dateInput(), ARRAY(DATE()));
  auto expected = dateAscNullSmallest();
  testSortArray(
      input,
      make_nullable_array_vector(expected, ARRAY(DATE())),
      make_nullable_array_vector(reverseNested(expected), ARRAY(DATE())));
}

TEST_F(SortArrayTest, bool) {
  auto input = make_nullable_array_vector(boolInput());
  auto expected = boolAscNullSmallest();
  testSortArray(
      input,
      make_nullable_array_vector(expected),
      make_nullable_array_vector(reverseNested(expected)));
}

TEST_F(SortArrayTest, array) {
  auto input = make_nullable_nested_array_vector(arrayInput());
  auto expected = arrayAscNullSmallest();
  testSortArray(
      input,
      make_nullable_nested_array_vector(expected),
      make_nullable_nested_array_vector(reverseNested(expected)));
}

TEST_F(SortArrayTest, map) {
  auto input = make_array_of_map_vector(mapInput());
  auto expected = mapAscNullSmallest();
  testSortArray(
      input,
      make_array_of_map_vector(expected),
      make_array_of_map_vector(reverseNested(expected)));
}

TEST_F(SortArrayTest, row) {
  auto rowType = ROW({INTEGER(), VARCHAR()});
  auto input = make_array_of_row_vector(rowType, rowInput());
  auto expected = rowAscNullSmallest();
  testSortArray(
      input,
      make_array_of_row_vector(rowType, expected),
      make_array_of_row_vector(rowType, reverseNested(expected)));
}
} // namespace
} // namespace kumo::pollux::functions::sparksql::test
